{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.autograd as autograd\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as Data\n",
    "torch.manual_seed(8) # for reproduce\n",
    "\n",
    "import time\n",
    "import numpy as np\n",
    "import gc\n",
    "import sys\n",
    "sys.setrecursionlimit(50000)\n",
    "import pickle\n",
    "torch.backends.cudnn.benchmark = True\n",
    "torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
    "from tensorboardX import SummaryWriter\n",
    "torch.nn.Module.dump_patches = True\n",
    "import copy\n",
    "import pandas as pd\n",
    "#then import my own modules\n",
    "from AttentiveFP import Fingerprint, Fingerprint_viz, save_smiles_dicts, get_smiles_dicts, get_smiles_array, moltosvg_highlight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit import Chem\n",
    "# from rdkit.Chem import AllChem\n",
    "from rdkit.Chem import QED\n",
    "from rdkit.Chem import rdMolDescriptors, MolSurf\n",
    "from rdkit.Chem.Draw import SimilarityMaps\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem\n",
    "from rdkit.Chem import rdDepictor\n",
    "from rdkit.Chem.Draw import rdMolDraw2D\n",
    "%matplotlib inline\n",
    "from numpy.polynomial.polynomial import polyfit\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm as cm\n",
    "import matplotlib\n",
    "import seaborn as sns; sns.set()\n",
    "from IPython.display import SVG, display\n",
    "import sascorer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of all smiles:  1128\n",
      "number of successfully processed smiles:  1128\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAU8AAAC/CAYAAAB+KF5fAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAAEx1JREFUeJzt3Xtsk9f9x/FPnIQQSIg7ZoG4bGsbHNWkJN0IFxWRjRLSTrRROjZKtZZobOmmRWNrworawTqEdiEZ0KVCdCBU6FUVaqGQCrY0KtNKxtAqNrpC7IJaoVaEAHFCSDFJ7N8f++HVs+PYB18S5/2SIuFzvn58fEg+fvzYz3nSfD6fTwCAqFiSPQAAGIkITwAwQHgCgAHCEwAMEJ4AYIDwBAADhCcAGCA8AcAA4QkABghPADBAeAKAAcITAAwQngBgICPZA4iFzs6r8npDLw41cWKOLl3qSfCIRi/mO/GY85tjsaTpllvGR32/lAhPr9c3aHje6EfiMN+Jx5wnHm/bAcAA4QkABghPADBAeAKAgZT4wGi46vdKnr7+sDVZmRnK4CUMGHEIzzjy9PXr+Kn2sDUld0xSRhb/DcBIwz4PABggPAHAAOEJAAYITwAwQHgCgAHCEwAMEJ4AYIDwBAADhCcAGCA8AcBAROcFnj9/Xjt37tS///1vnT59Wr29vdqzZ4/mzp0bULdo0SJ98sknQff/wQ9+oLq6uoC2ixcvqr6+Xu+88448Ho8cDofq6ur01a9+9SaeDgAkRkTh+fHHH6upqUkOh0Pz5s1TS0vLoLUlJSVBQTlp0qSA2x6PR1VVVert7dW6detktVq1e/duVVVV6dVXX5XD4TB4KgCQOBGFZ0lJiVpbWyVJzc3NYcNzwoQJKi4uDru9vXv3yuVy6fXXX9fMmTMlSXPmzNF9992nzZs3a+fOnZGOHwCSIqJjnhZLbA+NNjc3y263+4NTksaMGaOlS5fq6NGj6unhYlYAhreYf2D0t7/9TXfddZcKCwt1//336+WXX5bPF3hxKpfLJbvdHnTfgoICDQwM6OzZs7EeFgDEVEwXkvz617+uwsJCTZ8+XW63W2+++aZ+9atf6aOPPtKTTz7pr3O73crLywu6/422zs7OWA4LAGIupuG5fv36gNtlZWWqra3VCy+8oJUrV2rq1Kn+vrS0tEG3E64vlIkTc8L222y5UW0vVnyXe5WbMzZszbhxWbJ9YVyCRpQYyZrv0Yw5T7y4L2FeWVmpgwcP6l//+pc/PK1Wq9xud1BtV1eXvz8aly71DHrdapstVx0dV6IcdWz0evp1peda+JpejzoGBhI0ovhL5nyPVsz5zbFY0obcAQt5vziMJYDX6/3PA33uQ6f8/Hw5nc6g2ra2NqWnp+u2226L97AA4KbEPTz3798vi8WiO++8099WVlYmp9OpU6dO+duuX7+upqYmzZ8/Xzk50b8KAEAiRfy2/dChQ5KkkydPSpKOHz+uzs5OZWdnq7S0VAcPHtTbb7+t0tJSTZ48WV1dXXrzzTfV3NysVatWacqUKf5tLVu2TC+99JJqampUW1urvLw87dmzRxcuXNDWrVtj/BQBIPbSfP/7PaJBFBQUhGyfOnWqWlpadOLECW3dulUffvih3G63MjMzVVBQoOXLl6uysjLofh0dHdq0aZOOHDniPz2ztrZWs2fPjvpJDNdjnlc9kV09c3wKXT2T42+Jx5zfHNNjnhGH53BGeA4f/CEnHnN+c4btB0YAkIoITwAwQHgCgIHUOdg2QqVZ0nTV0x+2JiszQxm8zAHDCuGZZJ6+Af3T2RG2puSOScpIoQ+VgFTA/gwAGCA8AcAA4QkABghPADBAeAKAAcITAAwQngBggPAEAAOEJwAYIDwBwADhCQAGCE8AMMBqEyMAKy8Bww/hOQKw8hIw/LCvAgAGCE8AMEB4AoABwhMADBCeAGCA8AQAA4QnABggPAHAAOEJAAYiCs/z589r48aNWrFihe666y4VFBTo2LFjIWsPHDigBx54QHfeeacWLlyohoYGeTyeoLqLFy/qiSee0Ny5c1VcXKyHH35Y77333s09mwTq90pXPf1hf7y+ZI8SQLxEdD7fxx9/rKamJjkcDs2bN08tLS0h6/bv36+f//znWrFihZ588kmdOXNGDQ0N+uSTT7RlyxZ/ncfjUVVVlXp7e7Vu3TpZrVbt3r1bVVVVevXVV+VwOGLz7OLI09ev46faw9YU2W0JGg2ARIsoPEtKStTa2ipJam5uDhmeAwMDqq+v16JFi/T0009LkubNm6fMzEytW7dOVVVVKioqkiTt3btXLpdLr7/+umbOnClJmjNnju677z5t3rxZO3fujMVzA4C4iehtu8UydNmJEyfU0dGhysrKgPb7779fmZmZOnz4sL+tublZdrvdH5ySNGbMGC1dulRHjx5VT09PpONHFCI51NDvTfYogZEhZsvwuFwuSdKMGTMC2rOzszV9+nR//43auXPnBm2joKBAAwMDOnv2rGbNmhWroeH/RXKogdWZgMjE7K/E7XZLkvLy8oL68vLy/P03agerk6TOzs6oHnvixJyw/TZbblTbi4Tvcq9yc8aGrcnMzEhYzbhxWbJ9YVzYmkjGHMl2hhKP+UZ4zHnixXwXIy0tLaL2weqG6gvl0qUeeQf5aNtmy1VHx5WotheJXk+/rvRcC1vT15e4mt5ejzoGBsLXRDDmSLYTTrzmG4Njzm+OxZI25A5YyPvFagBWq1WSAvYwb+jq6grY07RarYPWfX5bADBcxSw88/PzJSng2KYkffbZZzp37lzAsdD8/Hw5nc6gbbS1tSk9PV233XZbrIY1aty4VAffOwUSI2Zv24uLi2Wz2bR//34tWbLE337w4EH19fUFtJWVlWnDhg06deqU7rjjDknS9evX1dTUpPnz5ysnJ/pd6NEukkt18L1TIHYiDs9Dhw5Jkk6ePClJOn78uDo7O5Wdna3S0lJlZGSotrZWa9eu1YYNG1ReXu7/knx5ebmKi4v921q2bJleeukl1dTUqLa2Vnl5edqzZ48uXLigrVu3xvgpAkDsRRyeq1evDrjd2NgoSZo6dar/S/OVlZWyWCzauXOnXnvtNd1yyy166KGH9JOf/CTgvllZWdq9e7c2bdqkp59+Wh6PRw6HQ7t27VJhYeHNPicAiLuIw7OtrS2iuoqKClVUVAxZZ7PZVF9fH+nDA8CwwqpKAGCA8AQAA4QnABggPAHAAOEJAAZYPgdR6/f+Z4WmUHyXe9Xr6VdWZoYyeGlGCiM8EbVwS9vl5ozVlZ5rLG2HlMdvNwLcOEc+HM6RBwhP/A/OkQciw1EpADBAeAKAAcITAAwQngBggPAEAAOEJwAYIDwBwADhCQAGCE8AMEB4AoABwhMADBCeAGCA8AQAA6yqhLgYamk7FkvGSEd4Ii6GWtqOxZIx0vHaDwAGCE8AMEB4AoCBmB50OnbsmB599NGQfW+99ZZuv/12/+13331XzzzzjE6fPq3x48errKxMdXV1mjBhQiyHBABxEZcj9nV1dSopKQlomzZtmv/fx44dU3V1te655x799Kc/1YULF9TQ0CCn06mXX35ZFgs7xACGt7iE56233qri4uJB++vr6zVjxgxt3brVH5Q2m03f+973dOjQIX3zm9+Mx7AAIGYSvovX3t6ukydPqqKiImAP8+6779akSZN0+PDhRA8JAKIWl/Bcv369HA6Hvva1r+mxxx7T+++/7+9zOp2SpBkzZgTdz263y+VyxWNIUen3Slc9/WF/uHY5MLrF9G17bm6uVq5cqTlz5shqterMmTP64x//qBUrVujFF19UUVGR3G63JCkvLy/o/nl5efrggw9iOSQjnr5+HT/VHraGa5cDo1tMw9PhcMjhcPhvz549W4sWLdLSpUu1ZcsWPf/88/6+tLS0kNsYrD2ciRNzwvbbbLlRbc93uVe5OWPD1mRmZlAziNycsUPWjBuXJdsXxoV9HEQu2t9x3Ly4nx9ns9m0YMECtbS0SJKsVqsk+fdAP6+rqyvkHulQLl3qkXeQ99E2W646Oq5Etb1eT7+u9FwLW9PXR00ouTljdaXn2pDb6e31qGNgIOzjIDImv+P4L4slbcgdsJD3i8NYgni9Xv+/bxzrDHVs0+l0hjwWCgDDTdzDs6OjQ0ePHvV/dWny5MkqLCzUgQMHAkK1tbVV7e3tWrJkSbyHBAA3LaZv22trazV9+nTNnDlTEyZM0NmzZ7Vjxw5du3ZNjz/+uL+urq5Oq1at0uOPP67ly5ervb1dDQ0NKioq0r333hvLIQFAXMQ0PAsKCtTU1KQXX3xRn332maxWq+bMmaMf/ehHstvt/rr58+dr+/btamxsVHV1tcaPH6/FixdrzZo1Sk9Pj+WQACAuYhqe1dXVqq6ujqh24cKFWrhwYSwfHgAShpPIAcAA4QkABrgOApJiqGscSVznCMMb4YmkGOoaRxLXOcLwxus6ABggPAHAAOEJAAYITwAwQHgCgAHCEwAMEJ4AYIDwBAADhCcAGCA8AcAA4QkABghPADDAqgtIef1eydPHCk6ILcITw1aslq3z9PXr+Kn2sDWs4IRo8duCYYtl6zCc8UYFAAwQngBggPAEAAOEJwAYIDwBwADhCQAGCE8AMMAX5DCiRfJFeq8vQYPBqEJ4YkSL5Iv0RXZbgkaD0SRp4Xn16lVt2bJFhw4dUnd3t/Lz8/XjH/9Y99xzT9wfe6hzndlTGX1idSooRo+khWdNTY0++OAD1dXVadq0aXrjjTdUU1Oj7du3q7S0NK6PPdS5zuypjD6xOhWURUhGj6SE55EjR3T06FE9++yzKisrkyTNmzdP586d029/+9u4hycQLyxCMnok5fXvz3/+s3JzcwPeoqelpamyslJnz57Vhx9+mIxhAUDEkvLy53K5lJ+fL4slMLsLCgokSU6nU/n5+RFvz2JJi6o/I92icWMzB60fqp+awWuyszI00J+ZmnOcmS5PvzdsjcUS2WMN9Ts74JWu9w+ErRmTka50y43HDb29obbz+W2MFNHOzVCG+r8YTJrP50v4xyPl5eX6yle+oueeey6g/aOPPlJ5ebl++ctf6uGHH070sAAgYkl7zUlLGzztw/UBwHCQlPC0Wq1yu91B7V1dXZKkvLy8RA8JAKKSlPDMz8/XmTNn5PUGHj9yOp2SJLvdnoxhAUDEkhKeZWVl6u7uVktLS0D7vn37dOutt0b1YREAJENSPm0vLS3V3Llz9dRTT8ntdmvatGnat2+f/vGPf2jbtm3JGBIARCUpn7ZLUk9PjzZv3qzDhw8HnJ65ePHiZAwHAKKStPAEgJFshH09FgCGB8ITAAykZHhevXpVGzdu1IIFCzRr1iw9+OCDevvtt5M9rBGvtbVVa9euVXl5uYqKirRw4ULV1NSora0tqPbdd9/Vd77zHc2aNUvz58/X+vXr1d3dnYRRp5bGxkYVFBSooqIiqI85T6yUDM+amhodOHBAq1ev1nPPPaf8/HzV1NToyJEjyR7aiPbKK6/o008/VVVVlXbs2KG1a9fq008/1bJly3TixAl/3bFjx1RdXa3Jkydr+/bteuKJJ9TS0qLq6uqg7/Yici6XSzt27NAXv/jFoD7mPAl8Keadd97x2e1235/+9Cd/m9fr9T300EO+e++9N4kjG/kuXrwY1NbV1eWbPXu2r6amxt/2rW99y1dRUeEbGBjwt/31r3/12e12X1NTU0LGmmoGBgZ83/72t30bNmzwffe73/U98MADAf3MeeKl3J4ny93Fz8SJE4PaJkyYoC9/+cs6f/68JKm9vV0nT55URUVFwKpZd999tyZNmqTDhw8nbLyp5Pnnn9f58+f1s5/9LKiPOU+OlAvPSJa7Q+xcvnxZLpdLM2bMkPTf+b1x+/PsdrtcLldCx5cKzp07pz/84Q9av369cnJygvqZ8+RIufB0u90hFxa50RZqQRKY8fl8Wrdunbxer1atWiXpv/M72P8B8x8dn8+nX/ziF1qwYMGgJ5Aw58mRktcCYLm7xNi0aZOam5v1m9/8RrfffntA32DzzPxH57XXXtP777+vt956a8ha5jyxUi48We4uMbZs2aJdu3bpqaee0oMPPuhvt1qtkkLv4Xd1dTH/Ubh8+bLq6+v12GOPKTs72/+1o/7+fnm9XnV3dysrK4s5T5KUe9vOcnfx98wzz2j79u1as2aNHn300YC+G8fdQh1nczqdIY/LIbT29nZduXJFv//971VSUuL/ee+99+R0OlVSUqLGxkbmPElSbs+zrKxMe/fuVUtLS8AxIpa7i41nn31W27Zt0+rVq/X9738/qH/y5MkqLCzUgQMHtHLlSv8Hd62trWpvb9eSJUsSPeQR60tf+pL27NkT1P7rX/9avb292rhxo6ZMmcKcJ0nKLQzi8/m0cuVKtbW1ac2aNf7l7vbt26dt27Zp0aJFyR7iiLVr1y797ne/0ze+8Q398Ic/DOgbM2aMHA6HpP/80a5atUpLlizR8uXL1d7eroaGBk2ZMkWvvPKK0tPTkzH8lPHII4+ou7tb+/fv97cx54mXcuEpsdxdvDzyyCP6+9//HrJv6tSpAYtb/+Uvf1FjY6NOnz6t8ePHa/HixVqzZg3H32IgVHhKzHmipWR4AkC8pdwHRgCQCIQnABggPAHAAOEJAAYITwAwQHgCgAHCEwAMEJ4AYIDwBAAD/wdyKveHcnpZxgAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 360x216 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "task_name = 'solubility'\n",
    "tasks = ['measured log solubility in mols per litre']\n",
    "\n",
    "raw_filename = \"../data/delaney-processed.csv\"\n",
    "feature_filename = raw_filename.replace('.csv','.pickle')\n",
    "filename = raw_filename.replace('.csv','')\n",
    "prefix_filename = raw_filename.split('/')[-1].replace('.csv','')\n",
    "smiles_tasks_df = pd.read_csv(raw_filename)\n",
    "smilesList = smiles_tasks_df.smiles.values\n",
    "print(\"number of all smiles: \",len(smilesList))\n",
    "atom_num_dist = []\n",
    "remained_smiles = []\n",
    "canonical_smiles_list = []\n",
    "for smiles in smilesList:\n",
    "    try:        \n",
    "        mol = Chem.MolFromSmiles(smiles)\n",
    "        atom_num_dist.append(len(mol.GetAtoms()))\n",
    "        remained_smiles.append(smiles)\n",
    "        canonical_smiles_list.append(Chem.MolToSmiles(Chem.MolFromSmiles(smiles), isomericSmiles=True))\n",
    "    except:\n",
    "        print(smiles)\n",
    "        pass\n",
    "print(\"number of successfully processed smiles: \", len(remained_smiles))\n",
    "smiles_tasks_df = smiles_tasks_df[smiles_tasks_df[\"smiles\"].isin(remained_smiles)]\n",
    "# print(smiles_tasks_df)\n",
    "smiles_tasks_df['cano_smiles'] =canonical_smiles_list\n",
    "assert canonical_smiles_list[8]==Chem.MolToSmiles(Chem.MolFromSmiles(smiles_tasks_df['cano_smiles'][8]), isomericSmiles=True)\n",
    "\n",
    "plt.figure(figsize=(5, 3))\n",
    "sns.set(font_scale=1.5)\n",
    "ax = sns.distplot(atom_num_dist, bins=28, kde=False)\n",
    "plt.tight_layout()\n",
    "# plt.savefig(\"atom_num_dist_\"+prefix_filename+\".png\",dpi=200)\n",
    "plt.show()\n",
    "plt.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_seed = 888 # 69，103, 107\n",
    "start_time = str(time.ctime()).replace(':','-').replace(' ','_')\n",
    "\n",
    "batch_size = 200\n",
    "epochs = 200\n",
    "\n",
    "p_dropout= 0.2\n",
    "fingerprint_dim = 200\n",
    "\n",
    "weight_decay = 5 # also known as l2_regularization_lambda\n",
    "learning_rate = 2.5\n",
    "radius = 2\n",
    "T = 2\n",
    "per_task_output_units_num = 1 # for regression model\n",
    "output_units_num = len(tasks) * per_task_output_units_num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Compound ID</th>\n",
       "      <th>ESOL predicted log solubility in mols per litre</th>\n",
       "      <th>Minimum Degree</th>\n",
       "      <th>Molecular Weight</th>\n",
       "      <th>Number of H-Bond Donors</th>\n",
       "      <th>Number of Rings</th>\n",
       "      <th>Number of Rotatable Bonds</th>\n",
       "      <th>Polar Surface Area</th>\n",
       "      <th>measured log solubility in mols per litre</th>\n",
       "      <th>smiles</th>\n",
       "      <th>cano_smiles</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>934</th>\n",
       "      <td>Methane</td>\n",
       "      <td>-0.636</td>\n",
       "      <td>0</td>\n",
       "      <td>16.043</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>-0.9</td>\n",
       "      <td>C</td>\n",
       "      <td>C</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Compound ID  ESOL predicted log solubility in mols per litre  \\\n",
       "934     Methane                                           -0.636   \n",
       "\n",
       "     Minimum Degree  Molecular Weight  Number of H-Bond Donors  \\\n",
       "934               0            16.043                        0   \n",
       "\n",
       "     Number of Rings  Number of Rotatable Bonds  Polar Surface Area  \\\n",
       "934                0                          0                 0.0   \n",
       "\n",
       "     measured log solubility in mols per litre smiles cano_smiles  \n",
       "934                                       -0.9      C           C  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if os.path.isfile(feature_filename):\n",
    "    feature_dicts = pickle.load(open(feature_filename, \"rb\" ))\n",
    "else:\n",
    "    feature_dicts = save_smiles_dicts(smilesList,filename)\n",
    "# feature_dicts = get_smiles_dicts(smilesList)\n",
    "remained_df = smiles_tasks_df[smiles_tasks_df[\"cano_smiles\"].isin(feature_dicts['smiles_to_atom_mask'].keys())]\n",
    "uncovered_df = smiles_tasks_df.drop(remained_df.index)\n",
    "uncovered_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_df = remained_df.sample(frac=0.2,random_state=random_seed)\n",
    "train_df = remained_df.drop(test_df.index)\n",
    "train_df = train_df.reset_index(drop=True)\n",
    "test_df = test_df.reset_index(drop=True)\n",
    "# print(len(test_df),sorted(test_df.cano_smiles.values))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "863604\n",
      "atom_fc.weight torch.Size([200, 39])\n",
      "atom_fc.bias torch.Size([200])\n",
      "neighbor_fc.weight torch.Size([200, 49])\n",
      "neighbor_fc.bias torch.Size([200])\n",
      "GRUCell.0.weight_ih torch.Size([600, 200])\n",
      "GRUCell.0.weight_hh torch.Size([600, 200])\n",
      "GRUCell.0.bias_ih torch.Size([600])\n",
      "GRUCell.0.bias_hh torch.Size([600])\n",
      "GRUCell.1.weight_ih torch.Size([600, 200])\n",
      "GRUCell.1.weight_hh torch.Size([600, 200])\n",
      "GRUCell.1.bias_ih torch.Size([600])\n",
      "GRUCell.1.bias_hh torch.Size([600])\n",
      "align.0.weight torch.Size([1, 400])\n",
      "align.0.bias torch.Size([1])\n",
      "align.1.weight torch.Size([1, 400])\n",
      "align.1.bias torch.Size([1])\n",
      "attend.0.weight torch.Size([200, 200])\n",
      "attend.0.bias torch.Size([200])\n",
      "attend.1.weight torch.Size([200, 200])\n",
      "attend.1.bias torch.Size([200])\n",
      "mol_GRUCell.weight_ih torch.Size([600, 200])\n",
      "mol_GRUCell.weight_hh torch.Size([600, 200])\n",
      "mol_GRUCell.bias_ih torch.Size([600])\n",
      "mol_GRUCell.bias_hh torch.Size([600])\n",
      "mol_align.weight torch.Size([1, 400])\n",
      "mol_align.bias torch.Size([1])\n",
      "mol_attend.weight torch.Size([200, 200])\n",
      "mol_attend.bias torch.Size([200])\n",
      "output.weight torch.Size([1, 200])\n",
      "output.bias torch.Size([1])\n"
     ]
    }
   ],
   "source": [
    "x_atom, x_bonds, x_atom_index, x_bond_index, x_mask, smiles_to_rdkit_list = get_smiles_array([canonical_smiles_list[0]],feature_dicts)\n",
    "num_atom_features = x_atom.shape[-1]\n",
    "num_bond_features = x_bonds.shape[-1]\n",
    "loss_function = nn.MSELoss()\n",
    "model = Fingerprint(radius, T, num_atom_features, num_bond_features,\n",
    "            fingerprint_dim, output_units_num, p_dropout)\n",
    "model.cuda()\n",
    "\n",
    "# optimizer = optim.Adam(model.parameters(), learning_rate, weight_decay=weight_decay)\n",
    "optimizer = optim.Adam(model.parameters(), 10**-learning_rate, weight_decay=10**-weight_decay)\n",
    "# optimizer = optim.SGD(model.parameters(), 10**-learning_rate, weight_decay=10**-weight_decay)\n",
    "\n",
    "tensorboard = SummaryWriter(log_dir=\"runs/\"+start_time+\"_\"+prefix_filename+\"_\"+str(fingerprint_dim)+\"_\"+str(p_dropout))\n",
    "\n",
    "model_parameters = filter(lambda p: p.requires_grad, model.parameters())\n",
    "params = sum([np.prod(p.size()) for p in model_parameters])\n",
    "print(params)\n",
    "for name, param in model.named_parameters():\n",
    "    if param.requires_grad:\n",
    "        print(name, param.data.shape)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, dataset, optimizer, loss_function):\n",
    "    model.train()\n",
    "    np.random.seed(epoch)\n",
    "    valList = np.arange(0,dataset.shape[0])\n",
    "    #shuffle them\n",
    "    np.random.shuffle(valList)\n",
    "    batch_list = []\n",
    "    for i in range(0, dataset.shape[0], batch_size):\n",
    "        batch = valList[i:i+batch_size]\n",
    "        batch_list.append(batch)   \n",
    "    for counter, batch in enumerate(batch_list):\n",
    "        batch_df = dataset.loc[batch,:]\n",
    "        smiles_list = batch_df.cano_smiles.values\n",
    "        y_val = batch_df[tasks[0]].values\n",
    "        \n",
    "        x_atom, x_bonds, x_atom_index, x_bond_index, x_mask, smiles_to_rdkit_list = get_smiles_array(smiles_list,feature_dicts)\n",
    "        atoms_prediction, mol_prediction = model(torch.Tensor(x_atom),torch.Tensor(x_bonds),torch.cuda.LongTensor(x_atom_index),torch.cuda.LongTensor(x_bond_index),torch.Tensor(x_mask))\n",
    "        \n",
    "        model.zero_grad()\n",
    "        loss = loss_function(mol_prediction, torch.Tensor(y_val).view(-1,1))     \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "def eval(model, dataset):\n",
    "    model.eval()\n",
    "    test_MAE_list = []\n",
    "    test_MSE_list = []\n",
    "    valList = np.arange(0,dataset.shape[0])\n",
    "    batch_list = []\n",
    "    for i in range(0, dataset.shape[0], batch_size):\n",
    "        batch = valList[i:i+batch_size]\n",
    "        batch_list.append(batch) \n",
    "    for counter, batch in enumerate(batch_list):\n",
    "        batch_df = dataset.loc[batch,:]\n",
    "        smiles_list = batch_df.cano_smiles.values\n",
    "#         print(batch_df)\n",
    "        y_val = batch_df[tasks[0]].values\n",
    "        \n",
    "        x_atom, x_bonds, x_atom_index, x_bond_index, x_mask, smiles_to_rdkit_list = get_smiles_array(smiles_list,feature_dicts)\n",
    "        atoms_prediction, mol_prediction = model(torch.Tensor(x_atom),torch.Tensor(x_bonds),torch.cuda.LongTensor(x_atom_index),torch.cuda.LongTensor(x_bond_index),torch.Tensor(x_mask))\n",
    "        MAE = F.l1_loss(mol_prediction, torch.Tensor(y_val).view(-1,1), reduction='none')        \n",
    "        MSE = F.mse_loss(mol_prediction, torch.Tensor(y_val).view(-1,1), reduction='none')\n",
    "#         print(x_mask[:2],atoms_prediction.shape, mol_prediction,MSE)\n",
    "        \n",
    "        test_MAE_list.extend(MAE.data.squeeze().cpu().numpy())\n",
    "        test_MSE_list.extend(MSE.data.squeeze().cpu().numpy())\n",
    "    return np.array(test_MAE_list).mean(), np.array(test_MSE_list).mean()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 14.101829 14.153732\n",
      "1 7.0114493 7.047961\n",
      "2 3.108138 2.8953424\n",
      "3 3.2585042 2.9606552\n",
      "4 3.05623 2.895416\n",
      "5 2.5793781 2.3981044\n",
      "6 2.37422 2.3506207\n",
      "7 2.051052 2.056384\n",
      "8 1.7632428 1.6644306\n",
      "9 1.2391969 1.2312653\n",
      "10 1.1727734 1.3196679\n",
      "11 1.1907371 1.3630211\n",
      "12 1.0583934 1.0375009\n",
      "13 0.87638783 0.9083488\n",
      "14 0.93653303 1.0779476\n",
      "15 0.7050104 0.7176549\n",
      "16 0.6185914 0.6122293\n",
      "17 0.5784847 0.54251933\n",
      "18 0.524155 0.4959931\n",
      "19 0.5165006 0.48837024\n",
      "20 0.49022672 0.49508625\n",
      "21 0.61975735 0.576389\n",
      "22 0.5594931 0.5415395\n",
      "23 0.4880609 0.4927354\n",
      "24 0.4781595 0.5386572\n",
      "25 0.48660624 0.55402905\n",
      "26 0.45836145 0.53724563\n",
      "27 0.4273047 0.44817877\n",
      "28 0.44057253 0.43573955\n",
      "29 0.39141765 0.3972685\n",
      "30 0.36541104 0.3871821\n",
      "31 0.3653211 0.42952332\n",
      "32 0.33889228 0.3858252\n",
      "33 0.34015888 0.4087695\n",
      "34 0.3191112 0.37919194\n",
      "35 0.34006742 0.42799726\n",
      "36 0.32399756 0.3658243\n",
      "37 0.30919543 0.36765808\n",
      "38 0.29378 0.3519741\n",
      "39 0.33591905 0.44943392\n",
      "40 0.45503053 0.59583575\n",
      "41 0.3982651 0.53883725\n",
      "42 0.36206883 0.5239284\n",
      "43 0.2861021 0.38699597\n",
      "44 0.26735854 0.37375143\n",
      "45 0.26806337 0.35091704\n",
      "46 0.25141892 0.34785795\n",
      "47 0.2455444 0.3234859\n",
      "48 0.24454862 0.3172578\n",
      "49 0.28436664 0.35613087\n",
      "50 0.23864563 0.33402294\n",
      "51 0.2165562 0.34272665\n",
      "52 0.20811984 0.32367724\n",
      "53 0.23260064 0.38025242\n",
      "54 0.2591624 0.43019146\n",
      "55 0.22236553 0.3492596\n",
      "56 0.22460997 0.29854274\n",
      "57 0.212699 0.30429307\n",
      "58 0.19937618 0.30977184\n",
      "59 0.20270747 0.3373879\n",
      "60 0.18915655 0.32384136\n",
      "61 0.19978929 0.35629952\n",
      "62 0.18274643 0.308038\n",
      "63 0.18363574 0.29763973\n",
      "64 0.18799883 0.32180622\n",
      "65 0.20422184 0.31658202\n",
      "66 0.17950775 0.30579573\n",
      "67 0.20967488 0.32631496\n",
      "68 0.16635548 0.29042864\n",
      "69 0.1612086 0.30518118\n",
      "70 0.18482341 0.3564476\n",
      "71 0.20123671 0.3971626\n",
      "72 0.16528417 0.3255045\n",
      "73 0.20383526 0.33607277\n",
      "74 0.16086927 0.32651183\n",
      "75 0.1706323 0.34154403\n",
      "76 0.15711367 0.30806687\n",
      "77 0.1631766 0.28714877\n",
      "78 0.16046616 0.28918526\n",
      "79 0.1702981 0.30676425\n",
      "80 0.14258052 0.30860266\n",
      "81 0.14483543 0.27927512\n",
      "82 0.13644794 0.29121304\n",
      "83 0.14682919 0.3242508\n",
      "84 0.19848785 0.38178527\n",
      "85 0.23597795 0.44870487\n",
      "86 0.19887677 0.37914205\n",
      "87 0.18625316 0.38788044\n",
      "88 0.14027187 0.3092822\n",
      "89 0.13562195 0.2820473\n",
      "90 0.16519098 0.28264424\n",
      "91 0.17290775 0.30866712\n",
      "92 0.17511961 0.29857272\n",
      "93 0.12964164 0.27815065\n",
      "94 0.13142663 0.2723677\n",
      "95 0.1267791 0.28297758\n",
      "96 0.12471318 0.27971077\n",
      "97 0.12716198 0.2916695\n",
      "98 0.1260275 0.29262406\n",
      "99 0.1177584 0.27966028\n",
      "100 0.14801198 0.36539343\n",
      "101 0.17111896 0.40786013\n",
      "102 0.16441946 0.38469306\n",
      "103 0.1597984 0.38099834\n",
      "104 0.14891054 0.35786965\n",
      "105 0.11926677 0.3216114\n",
      "106 0.10408543 0.29492736\n",
      "107 0.10954825 0.3092348\n",
      "108 0.112072386 0.32264322\n",
      "109 0.10044942 0.28704813\n",
      "110 0.109701425 0.30845708\n",
      "111 0.10291165 0.30223313\n",
      "112 0.09546712 0.28028372\n",
      "113 0.09166078 0.2713295\n",
      "114 0.09167679 0.284349\n",
      "115 0.09917971 0.3053903\n",
      "116 0.1237739 0.339058\n",
      "117 0.108224705 0.32561272\n",
      "118 0.116759524 0.32038623\n",
      "119 0.10574534 0.2928218\n",
      "120 0.11563 0.32742906\n",
      "121 0.12229207 0.34231102\n",
      "122 0.102857225 0.3116641\n",
      "123 0.10012016 0.2927059\n",
      "124 0.09488759 0.28436777\n",
      "125 0.09534983 0.28731188\n",
      "126 0.09602696 0.31940165\n",
      "127 0.087819315 0.29465392\n",
      "128 0.08088672 0.2795249\n",
      "129 0.08162902 0.2795669\n",
      "130 0.09592511 0.2999338\n",
      "131 0.08305369 0.26362723\n",
      "132 0.07901713 0.2581511\n",
      "133 0.07607882 0.2601807\n",
      "134 0.095511995 0.32585293\n",
      "135 0.07505002 0.26946563\n",
      "136 0.07576872 0.26892978\n",
      "137 0.073937446 0.2783528\n",
      "138 0.06976373 0.27668104\n",
      "139 0.07244427 0.28134292\n",
      "140 0.06801575 0.27249977\n",
      "141 0.0651211 0.26961178\n",
      "142 0.07776069 0.3155045\n",
      "143 0.09009632 0.33166087\n",
      "144 0.07395815 0.28318322\n",
      "145 0.066002846 0.2787307\n",
      "146 0.07096752 0.26585796\n",
      "147 0.063706875 0.2544428\n",
      "148 0.06787005 0.26727605\n",
      "149 0.08177305 0.293747\n",
      "150 0.090425394 0.26837406\n",
      "151 0.067982845 0.2744532\n",
      "152 0.06480571 0.28660515\n",
      "153 0.07299697 0.31155342\n",
      "154 0.08132879 0.3101688\n",
      "155 0.06463788 0.27831548\n",
      "156 0.06388941 0.26262227\n",
      "157 0.06615009 0.2862698\n",
      "158 0.076724656 0.32582256\n",
      "159 0.07614406 0.30823296\n",
      "160 0.06726112 0.26495948\n",
      "161 0.065908454 0.25958803\n",
      "162 0.07042457 0.26681674\n",
      "163 0.065994576 0.24290827\n",
      "164 0.063068524 0.25985658\n",
      "165 0.07012598 0.25770688\n",
      "166 0.08576999 0.24864702\n",
      "167 0.06799773 0.28267947\n",
      "168 0.06454735 0.25425777\n",
      "169 0.06362362 0.2720564\n",
      "170 0.07042562 0.26500762\n",
      "171 0.0557181 0.25662965\n",
      "172 0.054605052 0.26005578\n",
      "173 0.06132872 0.2704906\n",
      "174 0.05601138 0.27481642\n",
      "175 0.058675706 0.25861326\n",
      "176 0.054309905 0.2641674\n",
      "177 0.09296643 0.35202283\n",
      "178 0.10617179 0.34761968\n",
      "179 0.1046082 0.34988362\n",
      "180 0.09079263 0.3285547\n",
      "181 0.06286845 0.2708914\n",
      "182 0.0556561 0.27474353\n",
      "183 0.053465735 0.281348\n",
      "184 0.047210094 0.26962787\n",
      "185 0.05347685 0.27107787\n",
      "186 0.047920965 0.276868\n",
      "187 0.051813398 0.28031304\n",
      "188 0.04942369 0.25755265\n",
      "189 0.05781661 0.27909857\n",
      "190 0.04818734 0.27302265\n",
      "191 0.04532987 0.25569692\n",
      "192 0.059192445 0.2586186\n",
      "193 0.07201876 0.29116178\n",
      "194 0.04762533 0.2791224\n",
      "195 0.045088377 0.27431902\n",
      "196 0.047702346 0.28936514\n",
      "197 0.050864823 0.28118593\n",
      "198 0.044779234 0.2745504\n",
      "199 0.04241691 0.2562219\n",
      "200 0.040425576 0.26737574\n",
      "201 0.045696598 0.2912639\n",
      "202 0.045572106 0.2926771\n",
      "203 0.05161962 0.26930207\n",
      "204 0.06020265 0.3261002\n",
      "205 0.057267457 0.31491464\n",
      "206 0.05786306 0.3235378\n",
      "207 0.038328573 0.27838704\n",
      "208 0.06518547 0.324093\n",
      "209 0.052611817 0.2869405\n",
      "210 0.044522904 0.28711706\n",
      "211 0.043234237 0.2855541\n",
      "212 0.060058888 0.29349118\n",
      "213 0.042551696 0.25598413\n",
      "214 0.04536837 0.276069\n",
      "215 0.042736802 0.27733803\n",
      "216 0.041699495 0.27431542\n",
      "217 0.05037683 0.26516858\n"
     ]
    }
   ],
   "source": [
    "best_param ={}\n",
    "best_param[\"train_epoch\"] = 0\n",
    "best_param[\"test_epoch\"] = 0\n",
    "best_param[\"train_MSE\"] = 9e8\n",
    "best_param[\"test_MSE\"] = 9e8\n",
    "\n",
    "for epoch in range(800):\n",
    "    train_MAE, train_MSE = eval(model, train_df)\n",
    "    test_MAE, test_MSE = eval(model, test_df)\n",
    "#     tensorboard.add_scalars('MAE',{'train_MAE':test_MAE, 'test_MAE':test_MSE}, epoch)\n",
    "#     tensorboard.add_scalars('MSE',{'train_MSE':test_MAE, 'test_MSE':test_MSE}, epoch)\n",
    "    if train_MSE < best_param[\"train_MSE\"]:\n",
    "        best_param[\"train_epoch\"] = epoch\n",
    "        best_param[\"train_MSE\"] = train_MSE\n",
    "    if test_MSE < best_param[\"test_MSE\"]:\n",
    "        best_param[\"test_epoch\"] = epoch\n",
    "        best_param[\"test_MSE\"] = test_MSE\n",
    "        if test_MSE < 0.35:\n",
    "             torch.save(model, 'saved_models/model_'+prefix_filename+'_'+start_time+'_'+str(epoch)+'.pt')\n",
    "    if (epoch - best_param[\"train_epoch\"] >10) and (epoch - best_param[\"test_epoch\"] >18):        \n",
    "        break\n",
    "    print(epoch, train_MSE, test_MSE)\n",
    "    \n",
    "    train(model, train_df, optimizer, loss_function)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best epoch: 163 \n",
      " test MSE: 0.24290827\n"
     ]
    }
   ],
   "source": [
    "# evaluate model\n",
    "best_model = torch.load('saved_models/model_'+prefix_filename+'_'+start_time+'_'+str(best_param[\"test_epoch\"])+'.pt')     \n",
    "\n",
    "best_model_dict = best_model.state_dict()\n",
    "best_model_wts = copy.deepcopy(best_model_dict)\n",
    "\n",
    "model.load_state_dict(best_model_wts)\n",
    "(best_model.align[0].weight == model.align[0].weight).all()\n",
    "test_MAE, test_MSE = eval(model, test_df)\n",
    "print(\"best epoch:\",best_param[\"test_epoch\"],\"\\n\",\"test MSE:\",test_MSE)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
